#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

#include <arpa/inet.h>

#include "dns_server_rr_files.h"
#include "dns_packet_parsing.h"
#include "dns_util.h"

/*
 * Main
 */
int main(int argc, char ** argv) {

	/* Arguments */
	char * rr_file_path;
	int server_port;

	/* Read arguments */
	if (argc != 3) {
		fprintf(stderr, "Usage: %s rr_file_path server_port\n", argv[0]);
		exit(EXIT_FAILURE);
	}

	rr_file_path = argv[1];
	server_port = atoi(argv[2]);

	/* --- */

	SOCKET server_socket = INVALID_SOCKET;
	SOCKET client_socket = INVALID_SOCKET;
	SOCKADDR_IN server_addr;
	socklen_t server_addr_len;
	SOCKADDR_IN client_addr;
	socklen_t client_addr_len;

	/* Read server type */
	char server_domain[DOMAIN_STR_LENGTH];
	int server_type = read_server_type_from_conf(rr_file_path, server_domain);

	/* Create server socket  */
	if ((server_socket = socket(AF_INET, SOCK_STREAM, 0)) == INVALID_SOCKET)
		exit_syserr("socket()", NULL);

	/* Init server address */
	server_addr_len = sizeof(server_addr);
	memset(&server_addr, 0, server_addr_len);
	server_addr.sin_family = AF_INET;
	server_addr.sin_port = htons(server_port);
	server_addr.sin_addr.s_addr = htonl(INADDR_ANY);

	/* Bind socket */
	if (bind(server_socket, (SOCKADDR *) &server_addr, server_addr_len) == SOCKET_ERROR)
		exit_syserr("bind()", &server_socket, NULL);

	/* Listen for one client */
	if (listen(server_socket, 1) == SOCKET_ERROR)
		exit_syserr("listen()", &server_socket, NULL);

	/* Run server loop */
	char buffer[BUF_SIZE];
	int msg_length;
	while (true) {
		/* Reset buffer */
		memset(buffer, 0, BUF_SIZE);

		/* Wait for a request */
		printf("\nWaiting...\n");

		/* Accept the connection */
		if ((client_socket = accept(server_socket, (SOCKADDR *) &client_addr, &client_addr_len)) == INVALID_SOCKET)
			exit_syserr("accept()", &server_socket, NULL);

		/* Read the message */
		if (recv(client_socket, buffer, BUF_SIZE, 0) < 0)
			exit_syserr("recv()", &server_socket, &client_socket, NULL);
		printf("Message received\n");

		/* Extract requested name */
		char qname[DOMAIN_STR_LENGTH];
		char * name;
		char * domain;
		if (qname_from_question(buffer, BUF_SIZE, qname) < 0) {
			fprintf(stderr, "Cannot read qname\n");
			continue;
		}
		printf("QName: %s\n", qname);

		// Get subdomain, domain or tld if the server
		// is AUTHORITATIVE, TLD or ROOT respectively
		name = strtok(qname, ".");
		if (server_type == TLD)
			name = &(name[strlen(name) + 1]);
		else if (server_type == ROOT) {
			name = strtok(NULL, ".");
			name = &(name[strlen(name) + 1]);
		}

		printf("Name: %s\n", name);

		/* Check domain */
		if (server_type != ROOT) {
			// Get domain or tld if the server
			// is AUTHORITATIVE or TLD respectively
			if (server_type == AUTHORITATIVE)
				domain = &(name[strlen(name) + 1]);
			else {
				strtok(NULL, ".");
				domain = strtok(NULL, ".");
				*(domain - 1) = '.';
			}

			printf("Domain: %s\n", domain);

			if (strcmp(domain, server_domain)) {
				fprintf(stderr, "Domain does not match: '%s' != '%s'\n", domain, server_domain);

				// Build and send error
				msg_length = build_dns_reply_error(buffer, 8);
				send(client_socket, buffer, msg_length, 0);
				continue;
			}
		}

		/* Look for a record */
		char rr_data[DOMAIN_STR_LENGTH];
		if (!get_rr_data_from_name(rr_file_path, name, rr_data)) {
			fprintf(stderr, "No record found for '%s'\n", name);

			// Build and send error
			msg_length = build_dns_reply_error(buffer, 3);
			send(client_socket, buffer, msg_length, 0);
			continue;
		}

		/* Record found, build answer */
		unsigned char ip_addr[4];
		if (server_type == AUTHORITATIVE) {
			printf("HOST: %s\n", rr_data);

			ip_addr[0] = atoi(strtok(rr_data, "."));
			ip_addr[1] = atoi(strtok(NULL, "."));
			ip_addr[2] = atoi(strtok(NULL, "."));
			ip_addr[3] = atoi(strtok(NULL, "."));

			msg_length = build_dns_answer(buffer, ip_addr);
		} else {
			printf("SERVER: %s\n", rr_data);

			char ns_ip[IPV4_STR_LENGTH];
			if (!get_rr_data_from_name(rr_file_path, rr_data, ns_ip)) {
				fprintf(stderr, "No record found for '%s'\n", rr_data);
				exit(EXIT_FAILURE);
			}
			ip_addr[0] = atoi(strtok(ns_ip, "."));
			ip_addr[1] = atoi(strtok(NULL, "."));
			ip_addr[2] = atoi(strtok(NULL, "."));
			ip_addr[3] = atoi(strtok(NULL, "."));

			msg_length = build_dns_additional(buffer, name, rr_data, ip_addr);
		}

		/* Send answer */
		send(client_socket, buffer, msg_length, 0);

		shutdown(client_socket, SHUT_RDWR);
		close_socket(&client_socket);
	}

	/* --- */

	close_socket(&server_socket);

	return EXIT_SUCCESS;
}
